package com.jeesite.modules.ocr.adapter;

import java.awt.Color;
import java.awt.Graphics2D;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import javax.imageio.ImageIO;

import com.benjaminwan.ocrlibrary.OcrResult;
import com.benjaminwan.ocrlibrary.Point;
import com.benjaminwan.ocrlibrary.TextBlock;
import com.beust.jcommander.internal.Lists;
import com.jeesite.common.collect.ListUtils;
import com.jeesite.common.config.Global;
import com.jeesite.common.io.FileUtils;
import com.jeesite.modules.file.entity.FileUpload;
import com.jeesite.modules.ocr.base.IOcrAlgorithmConfig;
import com.jeesite.modules.ocr.base.IOcrResultAdapter;
import com.jeesite.modules.ocr.entity.OcrPoint;
import com.jeesite.modules.ocr.entity.OcrTempateConfigs;
import com.jeesite.modules.ocr.entity.OcrTextBlock;
import com.jeesite.modules.ocr.utils.OCRImageUtils;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.IoUtil;
import io.github.mymonstercat.ocr.InferenceEngine;
import io.github.mymonstercat.ocr.config.ParamConfig;

public class PaddleOcrResultAdapter implements IOcrResultAdapter {

	@Override
	public List<OcrTextBlock> process(FileUpload fileUpload, IOcrAlgorithmConfig ocrAlgorithmConfig)
			throws IOException {
		String originalImageFilePath = Global.getUserfilesBaseDir(fileUpload.getFileUrl());
		ParamConfig paramConfig = ParamConfig.getDefaultConfig();
		paramConfig.setDoAngle(true);
		paramConfig.setMostAngle(true);
		InferenceEngine engine = InferenceEngine.getInstance(io.github.mymonstercat.Model.ONNX_PPOCR_V3);
		OcrResult ocrResult = engine.runOcr(originalImageFilePath, paramConfig);
		// 读取源图片
		File originalImageFile = new File(originalImageFilePath);
		BufferedImage originalImage = ImageIO.read(originalImageFile);
		// 获取图片的宽度和高度
		int width = originalImage.getWidth();
		int height = originalImage.getHeight();
		// 创建一个同样大小的新图片
		BufferedImage blankImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
		// 填充整个图片为白色
		Graphics2D g2d = blankImage.createGraphics();
		g2d.setColor(Color.WHITE);
		g2d.fillRect(0, 0, width, height);
		g2d.dispose();

		// 将含有空白的北京图片 生成给前端。
		List<TextBlock> orcTextBlocks = ocrResult.getTextBlocks();
		// 根据配置信息，帅选掉 无关的textblocks
		List<TextBlock> needTextBlocks = Lists.newArrayList();

		for (TextBlock t : orcTextBlocks) {
			// 字符阈值
			if (ocrAlgorithmConfig.getOcrCharScores() == -1) {
				// 不限制
			} else {
				// 限制
				boolean charScoreOk = true;
				for (float d : t.getCharScores()) {
					if (d < ocrAlgorithmConfig.getOcrCharScores()) {
						charScoreOk = false;
					}
				}
				if (!charScoreOk) {
					continue;
				}
			}

			// 边框阈值
			if (ocrAlgorithmConfig.getOcrBoxScore() == -1) {
				// 不限制
			} else {
				// 限制
				if (t.getBoxScore() < ocrAlgorithmConfig.getOcrBoxScore()) {
					continue;
				}
			}
			// 角度阈值

			if (ocrAlgorithmConfig.getOcrAngleScore() == -1) {
				// 不限制
			} else {
				if (t.getAngleScore() < ocrAlgorithmConfig.getOcrAngleScore()) {
					continue;
				}
				// 限制
			}
			needTextBlocks.add(t);
		}

		// 保存新的空白图片
		String tempateImageRec = FileUtils.getFileNameWithoutExtension(originalImageFilePath) + "-rec."
				+ FileUtils.getFileExtension(fileUpload.getFileName());
		File outputImageFile = new File(tempateImageRec);
		ImageIO.write(blankImage, FileUtils.getFileExtension(fileUpload.getFileName()), outputImageFile);
		// FileUtil.copy(imgPath, drawPath, Boolean.TRUE);
		ByteArrayInputStream in = IoUtil.toStream(OCRImageUtils.drawImg(outputImageFile, needTextBlocks));
		FileUtil.writeFromStream(in, outputImageFile);

		List<OcrTextBlock> listOcrTextBlock = Lists.newArrayList();
		for (TextBlock t : needTextBlocks) {
			ArrayList<Point> orcPoints=t.getBoxPoint();
			ArrayList<OcrPoint> points = new ArrayList<>();
			for (Point point:orcPoints) {
				points.add(new OcrPoint(point.getX(),point.getY()));
			}
			listOcrTextBlock.add(new OcrTextBlock(points, t.getBoxScore(), t.getAngleIndex(), t.getAngleScore(),
					t.getAngleTime(), t.getText(), t.getCharScores(), t.getCrnnTime(), t.getBlockTime()));
		}
		return listOcrTextBlock;
	}

}
